#!/usr/bin/python
# -*- coding:utf-8 -*-
"""
给用户套餐增加主机模块
"""

import arrow
import datetime
from pymongo import MongoClient

from conf.settings import (
    CLOUDMONITOR_MONGODB,
    CLOUDMONITOR_DB,
)
from utility.log import Logger
logger = Logger(filename='user_add_host_vul.log').get_logger()

db = MongoClient(CLOUDMONITOR_MONGODB)[CLOUDMONITOR_DB]

# 需要开通主机套餐的用户名
users = []
if not users:
    logger.info("没有用户数据......")
    exit(-1)

# 过期时间, 默认三十天后
now = arrow.utcnow()
datetime_now = now.datetime
# expired_time = (datetime_now + datetime.timedelta(days=30)).replace(hour=16, minute=0, second=0, microsecond=0)

# 下次重置每月单次任务次数时间
next_reset_time = now.replace(
            day=1, hour=2, minute=0, second=0, microsecond=0).shift(months=1, hours=-8, minutes=-5).int_timestamp

#  主机资产数量
host_job_count = 100
# 每月单次任务次数
manual_task_count = 50
# 周期任务频率 （单位毫秒）, 默认一周一次
interval = 604800000


def main():
    db_user = db["users"]
    db_balance = db["balances"]
    db_charges = db["charges"]
    db_packages = db["packages"]
    free_package = db_packages.find_one({"key": "scanv_free"})
    if not free_package:
        logger.info("系统套餐数据错误")
        exit(-1)
    free_package_id = free_package["_id"]
    for username in users:
        logger.info(f"开始更新用户 {username} 数据")
        user = db_user.find_one({"username": username, "planId": {"$nin": [None, "", free_package_id]}})
        if not user:
            logger.info({f"用户: {username} 不存在或者没有购买套餐"})
            continue
        if not db_packages.find_one({"_id": user["planId"]}):
            logger.info({f"用户: {username} 套餐不存在(package 数据不存在)"})
            continue
        if not (balance := db_balance.find_one({"_id": user["balanceId"]})):
            logger.info({f"用户: {username} 套餐不存在(balance 数据不存在)"})
            continue
        if balance.get("feature", {}).get("hostVulTask"):
            logger.info({f"用户: {username} 已经开通了主机套餐"})
            continue

        expired_time = balance.get("planExpireTime")
        if not expired_time:
            logger.info({f"用户: {username} 到期时间异常: {expired_time}"})
            continue

        balance_update = {
            "$set": {
                "feature.hostVulTask": True, "feature.hostVulManualTask": True,
                "balance.host": {
                    "used": 0,
                    "total": host_job_count,
                    "number": 0,
                    "time": 0,
                    "taskType": "",
                    "limitType": "limit"
                },
                "balance.hostVulManualTask": {
                    "used": 0,
                    "total": manual_task_count,
                    "number": 0,
                    "time": 0,
                    "taskType": "",
                    "limitType": "period"
                },
                "balance.hostVulMaxFreq": {
                    "used": 0,
                    "total": interval,
                    "number": 0,
                    "time": 0,
                    "taskType": "",
                    "limitType": "frequency"
                }
            }
        }
        charges_data_host_vul_manual = {
            "limitType": "period",
            "used": 0,
            "total": manual_task_count,
            "number": 1,
            "funcType": "hostVulManualTask",
            "packId": user["planId"],
            "uid": user["_id"],
            "chargeType": "plan",
            "time": next_reset_time,
            "createTime": datetime_now,
            "updateTime": datetime_now,
            "expireTime": expired_time,
            "taskType": "hostVul",
            "isDelete": False
        }
        charges_data_host_max_freq = {
            "limitType": "frequency",
            "used": 0,
            "total": interval,
            "number": 1,
            "funcType": "hostVulMaxFreq",
            "packId": user["planId"],
            "uid": user["_id"],
            "chargeType": "plan",
            "time": 0,
            "createTime": datetime_now,
            "updateTime": datetime_now,
            "expireTime": expired_time,
            "taskType": "",
            "isDelete": False
        }
        charges_data_host = {
            "limitType": "limit",
            "used": 0,
            "total": host_job_count,
            "number": 1,
            "funcType": "host",
            "packId": user["planId"],
            "uid": user["_id"],
            "chargeType": "plan",
            "time": 0,
            "createTime": datetime_now,
            "updateTime": datetime_now,
            "expireTime": expired_time,
            "taskType": "",
            "isDelete": False
        }
        result = db_charges.insert_many([charges_data_host_vul_manual, charges_data_host_max_freq, charges_data_host])
        charges_ids = result.inserted_ids
        logger.info(f"用户 {username} charges写入成功， {charges_ids}")
        try:
            db_balance.update_one({"uid": user["_id"]}, balance_update)
            logger.info(f"用户 {username} balance更新成功")
        except:
            db_charges.remove({"_id": {"$in": charges_ids}})
            logger.info(f"用户 {username} balance更新失败， 删除charges数据")


if __name__ == "__main__":
    main()
